#!/bin/bash
# scand 2.0 
# written by zhongcheng@whu.edu.cn
# contact me if you have any suggestings or find bugs 
# QQ:32589927

function helpfunc() {
    echo 'scand 2.0: generate scan input file for gaussian/orca, requires a xyz file, gentor, xtb and xyz2QC program'
    echo 'usage: scand -d 1-2-3-4 -g "e10 0 180" abc.xyz; # use gentor to generate initial structure'
    echo '       scand -d 1-2-3-4 -s "0,180,19" abc.xyz; # use xtb to generate initial structure to scan dihedral 1-2-3-4'
    echo '       scand -b 2-3 -s "1.4,3.0,16" abc.xyz; # use xtb to generate initial structure to scan bond 2-3'
    echo '       note that for xtb scan the init value should be close to the input structure to avoid structure destroy'
    echo '       scand DOFT-SCANT-10-9-12-16 # extract energies in DOFT-SCANT-10-9-12-16'
    echo '       scand -E DOFT-SCANT-10-9-12-16 # extract energies in DOFT-SCANT-10-9-12-16 using excited state energy'
    echo '       -[d dihedral to scan, 1-2-3-4 for relaxed scan, 2-3 for rigid scan (gentor only)]'
    echo '       -[a angle to scan, 1-2-3 to specify the angle. e.g. scand -a 1-2-3 -s "90,180,9" abc.xyz ]'
    echo '       -[b bond  to scan, 1-2 to specify the bond. e.g. scand -b 1-2 -s "1.5,3.0,10" abc.xyz]'
    echo '       -[g scan dihedral setting for gentor, default is e10 0 180.'
    echo '           note this do not set absolute angle. so if initial angle is 3, the final angle is 3 13 23...'
    echo '       -[s scan setting for xtb, the format is -s "init,final,step", the step size is (final-init)/(step-1) ]'
    echo '       -[N OMP_NUM_THREADS used by xtb]' 
    echo '       -[F force constant for xtb scan, default is 50 ]'
    echo '       -[M maxcycle for xtb opt, default is 0, means maxstep ]'
    echo '       -[L level for xtb opt, default is crude, avail: crude,sloppy,loose,lax,normal,tight,vtight,extreme]'
    echo '       -[n %nproc for gaussian input file, default is 16]' 
    echo '       -[c charge and multplicity setting,  default is 0 1]'
    echo '       -[m %mem for gaussian input file, default is 4GB]' 
    echo '       -[k #p keywords for gaussian input file, default is opt(modredundant) pbe1pbe def2svp em(gd3bj) ]' 
    echo '       -[E using excited state energy]'
    echo '       -[p gau or orca ]' 
    echo '       -[S submit command that support the format such as g16s *.gjf ]'

    exit 1
}
if [ $# -lt 1 ]; then
    helpfunc
fi
while getopts ":n:N:g:s:S:d:a:b:F:c:m:k:p:M:L:E" opt
do
    case $opt in
        n) NPROC=$OPTARG;;
        N) NTHREAD=$OPTARG;;
        g) GENTOR=$OPTARG;;
        s) SCAN=$OPTARG;;
        d) dihedral=$OPTARG;;
        a) angle=$OPTARG;;
        b) bond=$OPTARG;;
        F) FORCECON=$OPTARG;;
        c) CHARGESPIN=$OPTARG;;
        m) MEM=$OPTARG;;
        k) KEYWORD=$OPTARG;;
        p) PROGRAM=$OPTARG;;
        M) MAXCYC=$OPTARG;;
        L) OPTLEVEL=$OPTARG;;
        S) SUBMIT_COM=$OPTARG;;
        E) EXCITED=1;;
        ?) helpfunc;;
    esac
done
shift $(($OPTIND - 1))

if [[ -z "$NPROC" ]];then NPROC=16;fi
if [[ -z "$CHARGESPIN" ]];then CHARGESPIN="0 1";fi
if [[ -z "$MEM" ]];then MEM="4GB";fi
if [[ -z "$PROGRAM" ]];then PROGRAM="gau";fi
if [[ -z "$FORCECON" ]];then FORCECON="50";fi
if [[ -z "$MAXCYC" ]];then MAXCYC="0";fi
if [[ -z "$OPTLEVEL" ]];then OPTLEVEL="loose";fi
#if [[ -z "$dihedral" && -z "$angle" && -z "$bond" ]];then
#    echo "You must input a bond/angle/dihedral !"
#    exit 1
#fi
if command -v gentor &> /dev/null; then gentor_com=gentor; else echo "gentor not available";fi
if command -v xyz2QC &> /dev/null; then xyz2QC_com=xyz2QC; else echo "xyz2QC not available, exit now!";exit 1; fi
if command -v xtb &> /dev/null; then xtb_com=xtb; else echo "xtb not available!"; fi
if command -v mfc.py &> /dev/null; then mfc_com=mfc; else echo "mfc not available!"; fi


if [[ -n "$xtb_com" && -z "$SCAN" ]];then 
    SCAN="0,180,19"
elif [[ -z "$xtb_com" && -n "$SCAN" ]];then 
    echo "Error! you set xtb scan with -s $SCAN but xtb command not available"
    exit 1
elif [[ -n "$gentor_com" && -z "GENTOR" ]];then
    GENTOR="e10 0 180"
elif [[ -z "$gentor_com" && -n "GENTOR" ]];then
    echo "Error! you set gentor scan with -g $GENTOR but gentor command not available"
    exit 1
elif [[ -z "$gentor_com" && -z "$xtb_com" ]];then
    echo "both xtb and gentor not found"
    exit 1
fi


function extract_ene() {
    BSNM=${1%%-SCAN*}
    cd $1
    logfiles=`ls ${BSNM}*.log`
    outfile=`ls ${BSNM}*.log`
    touch dft.exinfo
    touch dft.energy
    if [[ -n "$logfiles" ]];then
        if [[ "$EXCITED" == "1"  ]];then
            printf "%-30s%10s%20s\n" "file" "status" "abs_ex_E(au)" > dft.energy
            printf "%20s%20s\n" "osc.str." "trans.ene.(eV)" > dft.exinfo
        else
            printf "%-30s%10s%20s\n" "file" "status" "abs_E(au)" > dft.energy
        fi
        for f in $logfiles;do
            nf=`grep "Normal t" $f`
            if [[ -n $nf ]];then
                status="Normal"
            else
                status="Error"
            fi
            if [[ "$EXCITED" == "1"  ]];then
                ene=`awk '/Total Energy/{ene=$5}END{print ene}'  $f`
                osc=`awk '/Excited State   1/{ene=$9;gsub(/^f=/,"",ene)}END{print ene}' $f`
                te=`awk '/Excited State   1/{ene=$5}END{print ene}' $f`
                printf "%20s%20s\n" $osc $te >> dft.exinfo
            else
               ene=`awk '/SCF Done/{ene=$5}END{print ene}' $f`
            fi
            printf "%-30s%10s%20s\n" ${f%.*} $status $ene >> dft.energy
        done
        if [[ "$EXCITED" == "1" ]];then
            awk 'BEGIN{print "rel_ex_E(kcal/mol)"} NR == FNR { if(NR == 2 || $3 < min) min = $3; next } { if(FNR > 1) printf "%.3f\n",($3 - min)*627.51 }' dft.energy dft.energy > dftrelene.SCAND
        else
            awk 'BEGIN{print "rel_E(kcal/mol)"} NR == FNR { if(NR == 2 || $3 < min) min = $3; next } { if(FNR > 1) printf "%.3f\n",($3 - min)*627.51 }' dft.energy dft.energy > dftrelene.SCAND
        fi
    fi
    # calc boltzmann ditr.
    awk -v temp="298.15" '
        BEGIN{Z=0;idx=1;print "Boltzm.weight"}
        {
            if(NR>1){
                weight[idx]=exp(-1*$1/(0.001985875*temp));
                z += weight[idx];
                idx++;
            }
        }
        END{
            for(i=1;i<idx;i++){
                prob=weight[i];
                printf "%.3E\n",prob
            }
        }'  dftrelene.SCAND > boltzmann.SCAND
    # recombine columns
    awk '{print $1}' dft.energy > filename.SCAND
    awk '{print $2}' dft.energy > dftstatus.SCAND
    awk '{print $3}' dft.energy > dftabsene.SCAND
    awk '{print $1}' dft.exinfo > oscillator.SCAND
    awk '{print $2}' dft.exinfo > transene.SCAND
    if [ -f "SCAN.DATA" ];then
        awk '{print $1}' SCAN.DATA > measure.SCAND
        awk '{print $2}' SCAN.DATA > xtbene.SCAND
        paste filename.SCAND measure.SCAND dftrelene.SCAND boltzmann.SCAND oscillator.SCAND transene.SCAND dftstatus.SCAND dftabsene.SCAND xtbene.SCAND > allinfo.SCAND 
    else
        paste filename.SCAND dftrelene.SCAND boltzmann.SCAND oscillator.SCAND transene.SCAND  dftstatus.SCAND dftabsene.SCAND > allinfo.SCAND 
    fi
    # format output
    awk '
        NR == FNR {
          for (i = 1; i <= NF; i++) {
            if (length($i) > col_width[i])
              col_width[i] = length($i)
          }
          next
        }
        FNR == 1 {
          for (i = 2; i <= length(col_width); i++)
            col_width[i] += 2
        }
        {
          for (i = 1; i <= NF; i++)
            printf "%*s ", col_width[i], $i
          print "" 
        }
        ' allinfo.SCAND allinfo.SCAND | tee SCAN_RESULT.DATA
    rm *.SCAND
}

# convert structure to xyz
if [[ "${1##*.}" == "xyz" ]];then
    BASENAME=${1/.xyz/}
    cp $1 mol.xyz
elif [[ -d $1 ]];then
    extract_ene $1
    exit 1
elif [[ -n "$mfc_com" ]];then
    echo "converting $1 to xyz file using $mfc_com"
    $mfc_com -f xyz $1
    BASENAME=${1%.*}
    cp ${BASENAME}.xyz mol.xyz
    if [ $? -ne 0 ];then
        echo "Input file format convert to xyz failed"
        exit 1
    fi
else
    echo "Could not transfer input format to xyz"
    exit 1
fi

# use xtb or gentor to generate xyz traj
if [[ -n "$GENTOR" ]];then
    if [[ $dihedral =~ ^[0-9]+-[0-9]+-[0-9]+-[0-9]+$ ]];then
        tail=${dihedral#*-}
        BOND=${tail%-*}
    elif [[ $dihedral =~ ^[0-9]+-[0-9]+$ ]];then
        BOND=${dihedral}
    else
        echo "You set -g $GENTOR but do not set -d option"
        exit 1
    fi
    if [[ -n "$gentor_com" ]];then
        wd=${BASENAME}-SCAND-$dihedral
        mkdir $wd
        cp mol.xyz $wd
        cd $wd
        echo SCAND-$dihedral > SCAN.DATA 
        echo "$GENTOR" | awk '{if($1 ~/^e/){
                gsub(/^e/,"",$1)
                ss=$1;
                if(NF==3){
                    init=$2;final=$3
                } else if(NF==1){
                    init=0;final=360
                };
                if(ss < 0){
                    ss=ss*-1
                    temp=init
                    init=final
                    final=temp
                };
                idx=1;
                for(i=init;i<=final;i+=ss){
                    arr[idx++]=i
                }
            } else
                 for(k=1;k<=NF;k++){
                    arr[k]=$k
                }
            }END{for(j in arr){print "+"arr[j]}}' >> SCAN.DATA
        echo "Entering $wd"
        echo "$BOND" > gentor.ini 
        echo "$GENTOR" >> gentor.ini
        echo "gentor: rotate $BOND by $GENTOR"
        echo "Running  echo -en \"\n\" | $gentor_com"
        echo -en "\n" | $gentor_com > gentor.out
        if [ $? -ne 0 ];then
            echo "gentor run error!!! You should run following commands to check what happens"
            echo 'echo -en  "\n" | gentor'
            exit 1
        fi
    else
        echo "You set -g $GENTOR but gentor command not found"
        exit 1
    fi
elif [[ -n "$SCAN" ]];then
    SCAN_STR=""
    if [[ $bond =~ ^[0-9]+-[0-9]+$ ]];then
        LBL=SCANB-$bond
        init=$(echo $SCAN | cut -d ',' -f1)
        bond=$(echo $bond | tr "-" ",")
        final=$(echo $SCAN | cut -d ',' -f2)
        step=$(echo $SCAN | cut -d ',' -f3)
        STEP_SIZE=`echo "scale=3;($final-$init)/($step-1)" | bc`
        SCAN_STR="distance: ${bond} with $SCAN, stepsize: $STEP_SIZE"
        SCAN_XTB="distance: ${bond},$init; $SCAN"
    elif [[ -n $bond ]];then
        echo "the format of -b $bond is not right"
        exit 1
    fi
    if [[ $angle =~ ^[0-9]+-[0-9]+-[0-9]+$ ]];then
        LBL=SCANA-$angle
        init=$(echo $SCAN | cut -d ',' -f1)
        angle=$(echo $angle | tr "-" ",")
        final=$(echo $SCAN | cut -d ',' -f2)
        step=$(echo $SCAN | cut -d ',' -f3)
        STEP_SIZE=`echo "scale=3;($final-$init)/($step-1)" | bc`
        SCAN_STR="angle: ${angle} with $SCAN, stepsize: $STEP_SIZE"
        SCAN_XTB="angle: ${angle},$init; $SCAN"
    elif [[ -n $angle ]];then
        echo "the format of -a $angle is not right"
        exit 1
    fi
    if [[ $dihedral =~ ^[0-9]+-[0-9]+-[0-9]+-[0-9]+$ ]];then
        LBL=SCANT-$dihedral
        init=$(echo $SCAN | cut -d ',' -f1)
        if [[ $init == 0 ]];then
            init=0.001
        fi
        dihedral=$(echo $dihedral | tr "-" ",")
        final=$(echo $SCAN | cut -d ',' -f2)
        if [[ $final == 0 ]];then
            final=0.001
        fi
        step=$(echo $SCAN | cut -d ',' -f3)
        STEP_SIZE=`echo "scale=3;($final-$init)/($step-1)" | bc`
        SCAN_STR="dihedral: ${dihedral} with $SCAN, stepsize: $STEP_SIZE"
        SCAN_XTB="dihedral: ${dihedral},$init; $SCAN"
    elif [[ -n $dihedral ]];then
        echo "the format of -d $dihedral is not right for xtb"
        exit 1
    fi
    if [[ -z "$SCAN_STR" ]];then
        echo "You set -s $SCAN but do not set valid -d or -a or -b option"
        exit 1
    else
        wd=${BASENAME}-${LBL}
        mkdir $wd
        cp mol.xyz $wd
        cd $wd
        echo $LBL > SCAN.DATA
        printf "%.3f\n" $init >> SCAN.DATA
        for i in `seq $((step-1))`;do
            echo  "scale=3;$init+$i*$STEP_SIZE" | bc >> SCAN.DATA
        done
        echo "Entering $wd"
        echo '$constrain' > xtb_scan.inp
        echo "  force constant=$FORCECON" >> xtb_scan.inp
        echo '$scan' >> xtb_scan.inp
        echo "  $SCAN_XTB" >> xtb_scan.inp
        echo '$end' >> xtb_scan.inp
        if [[ -n $MAXCYC ]];then
            echo '$opt' >> xtb_scan.inp
            echo "  maxcycle=$MAXCYC" >> xtb_scan.inp
            echo '$end' >> xtb_scan.inp
        fi
        echo "xtb: scan $SCAN_STR, constrain:$FORCECON, level:$OPTLEVEL, maxcycle:$MAXCYC"
        charge=$(echo "$CHARGESPIN" | awk '{print $1}')
        multi=$(echo "$CHARGESPIN" | awk '{print $2}')
        spin=`echo "$multi - 1" | bc`
        if [[ -n "$NTHREAD" ]];then
            export OMP_NUM_THREADS=$NTHREAD
        fi
        echo "running xtb mol.xyz --opt $OPTLEVEL -c $charge -u  $spin --input xtb_scan.inp"
        xtb mol.xyz --opt $OPTLEVEL -c $charge -u  $spin --input xtb_scan.inp  | tee xtb_scan.out  | awk '{if($0 ~ /... step/)printf "%s\r",$0}'
        if [ $? -ne 0 ];then
            echo "xtb run error!!! You run following commands to check what happens:"
            echo "xtb mol.xyz --opt -c $charge -u  $spin --input xtb_scan.inp"
            exit 1
        else
            mv xtbscan.log traj.xyz 
            echo "xtb energy:"
            awk '/energy/{if(E0==0)E0=$2;i=i+1;printf "step %02d: %8.3f kCal/mol\n",i,(E0-$2)*627.51}' traj.xyz | tee xtb.energy
            awk 'BEGIN{print "xtb_ene(kCal/mol)"} NR == FNR { if(NR == 1 || $3 < min) min = $3; next } { printf "%.3f\n",$3 - min }' xtb.energy xtb.energy > tmp.energy
            paste SCAN.DATA tmp.energy | awk '{printf "%-20s%20s\n", $1, $2}' > scand.tmp
            mv scand.tmp SCAN.DATA
            rm tmp.energy
        fi
    fi
fi

if [[ "${PROGRAM,,}" =~ "gau" ]];then 
    if [[ -z "$KEYWORD" ]];then KEYWORD="pbe1pbe def2svp em(gd3bj) g09default";fi
    if [[ $dihedral =~ ^[0-9]+[-,][0-9]+[-,][0-9]+[-,][0-9]+$ ]];then
        FITD="D ${dihedral//-/ } F"
        if [[ ! ${KEYWORD,,} =~ opt[=\(]+.*modr ]];then
            KEYWORD="${KEYWORD} opt(modredundant)"
            echo $KEYWORD
        fi
    elif [[ $dihedral =~ ^[0-9]+[-,][0-9]+$ ]];then
        FITD=" "
    fi  
    echo "%nproc=$NPROC" > template.gjf
    echo "%mem=$MEM"     >> template.gjf
    echo "#p $KEYWORD"   >> template.gjf
    echo -en "\nTemplate\n\n" >> template.gjf
    echo "$CHARGESPIN"   >> template.gjf
    echo "[GEOMETRY]"    >> template.gjf
    echo -en "\n$FITD\n\n\n" >> template.gjf
    echo 2 > xyz2qc.inp
    echo traj.xyz >> xyz2qc.inp
    echo -en "\n${BASENAME}_\n\n\n" >> xyz2qc.inp
    INPFILE="${BASAME}*.gjf"
    

elif [[ "${PROGRAM,,}" =~ "orca" ]];then
    if [[ -z "$KEYWORD" ]];then KEYWORD="PBE0 def2-SVP D3BJ";fi
    if [[ $dihedral =~ ^[0-9]+[-,][0-9]+[-,][0-9]+[-,][0-9]+$ ]];then
        FITD="%geom Constraints {D ${dihedral//-/ } C} end end"
        if [[ ! ${KEYWORD,,} =~ opt ]];then
            KEYWORD="${KEYWORD} opt"
        fi
    elif [[ $dihedral =~ ^[0-9]+[-,][0-9]+$ ]];then
        FITD=" "
    fi  
    meml=${MEM,,}
    memv=${meml/gb/}
    MC=`echo "$memv * 1024 / $NPROC" | bc`
    echo "! $KEYWORD" > template.inp
    echo "%pal nprocs $NPROC" >> template.inp
    echo "end"                >> template.inp
    echo "%MaxCore $MC"       >> template.inp
    echo "$FITD"              >> template.inp
    echo "* xyz $CHARGESPIN"  >> template.inp
    echo "[GEOMETRY]"         >> template.inp
    echo "*"                  >> template.inp
    echo 4 > xyz2qc.inp
    echo traj.xyz >> xyz2qc.inp
    echo -en "\n${BASENAME}_\n\n\n" >> xyz2qc.inp
    INPFILE="${BASAME}*.inp"
fi

cat xyz2qc.inp | xyz2QC > xyz2QC.out
exitcode=$?
if [ $exitcode -ne 0 ];then
    echo "xyz2QC run error!!! You should check output of xyz2QC"
    exit
else
    echo -n "xyz2QC:"
    grep Done xyz2QC.out | sed "s/current folder/$wd/"
    rm template.gjf template.inp 2> /dev/null
    if [[ -n "$SUBMIT_COM" ]] ;then
        echo "$SUBMIT_COM $INPFILE"
        $SUBMIT_COM $INPFILE
    fi
fi



